from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 基础数据配置 ==========
# 肿瘤相关诊断
CANCER_DIAGNOSES = [
    ("C34.901", "肺恶性肿瘤", "T2N1M0"),
    ("C50.902", "乳腺恶性肿瘤", "T1N0M0"),
    ("C16.901", "胃恶性肿瘤", "T3N2M0"),
    ("C18.901", "结肠恶性肿瘤", "T4N1M1"),
    ("C22.001", "肝细胞癌", "T2N0M0")
]

# 医院信息
HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]

# 生命体征正常范围
VITAL_SIGNS = {
    'temperature': (36.5, 37.5),  # 体温(℃)
    'pulse': (60, 100),  # 心率(次/分)
    'respiratory': (16, 20),  # 呼吸(次/分)
    'bp_systolic': (90, 140),  # 收缩压(mmHg)
    'bp_diastolic': (60, 90)  # 舒张压(mmHg)
}

# ECOG/KPS评分
PERFORMANCE_STATUS = {
    'ECOG': ["0", "1", "2", "3", "4"],
    'KPS': ["100", "90", "80", "70", "60", "50", "40", "30", "20", "10"]
}


# ========== 时间范围配置 ==========
def generate_day_range(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_day_range("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 12 if date_obj.day % 2 == 0 else 7


# ========== 核心生成函数 ==========
def generate_admission_records():
    sql_statements = []
    record_count = 0

    # 初始化计数器
    counters = {
        'patient_id': 2001,
        'visit_sn': 2001,
        'record_sn': 10001,
        'inpatient_no': 500001
    }

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1

            # 患者基础信息
            patient_id = f"PT{counters['patient_id']}"
            counters['patient_id'] += 1

            visit_sn = f"VIS{counters['visit_sn']}"
            counters['visit_sn'] += 1

            record_sn = f"REC{counters['record_sn']}"
            counters['record_sn'] += 1

            # 医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 基础医疗信息
            medical_record_no = f"MR{random.randint(100000, 999999)}"
            inpatient_no = f"IP{counters['inpatient_no']}"
            counters['inpatient_no'] += 1

            # 诊断信息
            diag_code, diag_name, tnm_staging = random.choice(CANCER_DIAGNOSES)
            t, n, m = tnm_staging[1], tnm_staging.split("N")[1].split("M")[0], tnm_staging[-1]

            # 时间信息
            record_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d"),
                end_date=datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1)
            ).strftime('%Y-%m-%d %H:%M:%S')

            # 生命体征
            temp = round(random.uniform(*VITAL_SIGNS['temperature']), 1)
            pulse = random.randint(*VITAL_SIGNS['pulse'])
            resp = random.randint(*VITAL_SIGNS['respiratory'])
            bp_sys = random.randint(*VITAL_SIGNS['bp_systolic'])
            bp_dia = random.randint(*VITAL_SIGNS['bp_diastolic'])

            # 体格检查
            height = random.randint(155, 185)
            weight = random.randint(45, 85)
            bsa = round(0.0061 * height + 0.0124 * weight - 0.0099, 2)  # 体表面积计算

            # 生成病历内容
            chief_complaint = f"发现{diag_name.replace('恶性肿瘤', '')}肿物{random.randint(1, 12)}个月"
            current_history = f"患者{random.randint(3, 24)}个月前无意中发现{diag_name.replace('恶性肿瘤', '')}肿物，大小约{random.randint(1, 5)}cm，"
            current_history += f"伴{random.choice(['疼痛', '消瘦', '食欲减退', '无特殊不适'])}，未予重视。"

            physical_exam = f"体温{temp}℃，脉搏{pulse}次/分，呼吸{resp}次/分，血压{bp_sys}/{bp_dia}mmHg\n"
            physical_exam += f"身高{height}cm，体重{weight}kg，体表面积{bsa}m²\n"
            physical_exam += f"专科检查：{random.choice(['可触及肿块', '局部压痛', '皮肤改变', '淋巴结肿大'])}"

            data = {
                'auxiliary_exam': f"影像学检查提示{diag_name}可能，建议病理检查确诊",
                'body_surface_area': str(bsa),
                'body_temperature': str(temp),
                'chief_complaint': chief_complaint,
                'confirmed_diagnosis': diag_name,
                'content_type': "XML",
                'current_medhistory': current_history,
                'diastolic_pressure': str(bp_dia),
                'ecog_score': random.choice(PERFORMANCE_STATUS['ECOG']),
                'family_history': f"家族中{random.choice(['有', '无'])}肿瘤病史",
                'from_table': 'EMR_ADMISSION',
                'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                'heart_rate': str(pulse),
                'height': str(height),
                'hospital_code': hospital_code,
                'hospital_name': hospital_name,
                'hospitalization_times': str(random.randint(1, 3)),
                'inpatient_no': inpatient_no,
                'kps_score': random.choice(PERFORMANCE_STATUS['KPS']),
                'marriage_birth_history': f"已婚，育有{random.randint(0, 2)}子",
                'medical_record_no': medical_record_no,
                'menstrual_history': random.choice(["月经规律", "已绝经", "月经不规律"]),
                'other_diagnosis': random.choice(["高血压", "糖尿病", "冠心病", "无"]),
                'past_medhistory': random.choice(["既往体健", "有手术史", "有慢性病史"]),
                'patient_id': patient_id,
                'patient_id_old': f"OLD_{patient_id}",
                'personal_medhistory': random.choice(["吸烟史", "饮酒史", "无特殊嗜好"]),
                'physical_exam': physical_exam,
                'primary_diagnosis': diag_name,
                'record_datetime': record_datetime,
                'record_sn': record_sn,
                'record_status': 1,
                'record_template_name': "肿瘤科入院记录模板",
                'record_text': f"<xml>{chief_complaint}</xml>",
                'record_title': f"{diag_name}入院记录",
                'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'respiratory_rate': str(resp),
                'revision_diagnosis': None,
                'signature_doctor': fake.name(),
                'special_exam': physical_exam.split("专科检查：")[1],
                'systolic_pressure': str(bp_sys),
                'visit_sn': visit_sn,
                'weight': str(weight),
                'yy_collection_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'yy_record_md5': f"md5_{random.getrandbits(128):032x}",
                'yy_upload_status': 0,
                'yy_etl_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'yy_upload_time': None,
                'yy_batch_time': day,
                'yy_record_batch_id': f"BATCH{day}_{i}",
                'yy_backfill_time': None,
                'yy_backfill_status': None,
                'branch_code': branch_code,
                'branch_name': branch_name,
                'date_for_partition': record_datetime
            }

            sql = _generate_sql('b05_1', data)
            sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_admission_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b05_1')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b05_1_admission_records.sql")